import matplotlib.pyplot as plt
import cubic_spline_planner
import numpy as np
from math import sin, cos, radians

dt = 0.1  # [s] time difference
L = 2.9  # [m] Wheel base of vehicle
max_steer = np.radians(30.0)  # [rad] max steering angle

#  target course
ax = [0.0, 100.0, 100.0, 50.0, 60.0]
ay = [0.0, 0.0, -30.0, -20.0, 0.0]
cx, cy, cyaw, ck, s = cubic_spline_planner.calc_spline_course(ax, ay, ds=0.1)
curve = np.array([cx, cy])

def normalize_angle(angle):
    """
    Normalize an angle to [-pi, pi].

    :param angle: (float)
    :return: (float) Angle in radian in [-pi, pi]
    """
    while angle > np.pi:
        angle -= 2.0 * np.pi

    while angle < -np.pi:
        angle += 2.0 * np.pi

    return angle

class Robot(object):
    """
    Class representing the state of a vehicle.

    :param x: (float) x-coordinate
    :param y: (float) y-coordinate
    :param yaw: (float) yaw angle
    :param v: (float) speed
    """

    def __init__(self, x=0.0, y=0.0, yaw=0.0, v=0.0):
        """Instantiate the object."""
        super(Robot, self).__init__()
        self.x = x
        self.y = y
        self.yaw = yaw
        self.v = v

    def point_in_local(self, p):
        ret = np.array([
            [cos(self.yaw), sin(self.yaw)],
            [-sin(self.yaw), cos(self.yaw)]
        ]).dot(np.array([
            p[0, :] - self.x,
            p[1, :] - self.y
        ]))
        return ret

    def update(self, acceleration, delta):
        """
        Update the state of the vehicle.

        Stanley Control uses bicycle model.

        :param acceleration: (float) Acceleration
        :param delta: (float) Steering
        """
        delta = np.clip(delta, -max_steer, max_steer)

        self.x += self.v * np.cos(self.yaw) * dt
        self.y += self.v * np.sin(self.yaw) * dt
        self.yaw += self.v / L * np.tan(delta) * dt
        self.yaw = normalize_angle(self.yaw)
        self.v += acceleration * dt

def calc_target_index(state:Robot, cx, cy):
    """
    Compute index in the trajectory list of the target.

    :param state: (State object)
    :param cx: [float]
    :param cy: [float]
    :return: (int, float)
    """
    # Calc front axle position
    fx = state.x + L * np.cos(state.yaw)
    fy = state.y + L * np.sin(state.yaw)

    # Search nearest point index
    dx = [fx - icx for icx in cx]
    dy = [fy - icy for icy in cy]
    d = np.hypot(dx, dy)
    target_idx = np.argmin(d)

    # Project RMS error onto front axle vector
    front_axle_vec = [-np.cos(state.yaw + np.pi / 2),
                      -np.sin(state.yaw + np.pi / 2)]
    error_front_axle = np.dot([dx[target_idx], dy[target_idx]], front_axle_vec)

    return target_idx, error_front_axle

def pid_control(d, robot_yaw, keep_yaw):
    vth = (normalize_angle(keep_yaw- robot_yaw) + d*0.8 ) * 1.5
    return vth

def speed_control(target, current):
    """
    Proportional control for the speed.

    :param target: (float)
    :param current: (float)
    :return: (float)
    """
    return 1.5 * (target - current)

def mapid_control(state, cx, cy, cyaw, last_target_idx):
    """
    Stanley steering control.

    :param state: (State object)
    :param cx: ([float])
    :param cy: ([float])
    :param cyaw: ([float])
    :param last_target_idx: (int)
    :return: (float, int)
    """
    current_target_idx, error_front_axle = calc_target_index(state, cx, cy)

    if last_target_idx >= current_target_idx:
        current_target_idx = last_target_idx

    # theta_e corrects the heading error
    local_point = state.point_in_local(np.array([
        [cx[current_target_idx]], 
        [cy[current_target_idx]]]
        ))
    delta = pid_control(local_point[1, 0], state.yaw, cyaw[current_target_idx])

    return delta, current_target_idx

state = Robot(x=-0.0, y=20.0, yaw=np.radians(20.0), v=0.0)
rp = state.point_in_local(curve) # 机器人坐标下的轨迹点

target_speed = 30.0 / 3.6  # [m/s]
max_simulation_time = 100.0
time = 0.0
target_idx, _ = calc_target_index(state, cx, cy)
last_idx = len(cx) - 1
x = [state.x]
y = [state.y]
yaw = [state.yaw]
v = [state.v]
t = [0.0]

while max_simulation_time >= time and last_idx > target_idx:
    di, target_idx = mapid_control(state, cx, cy, cyaw, target_idx)
    print(cyaw[target_idx])
    ai = speed_control(target_speed, state.v)
    state.update(ai, di)
    time += dt

    x.append(state.x)
    y.append(state.y)
    yaw.append(state.yaw)
    v.append(state.v)
    t.append(time)

    plt.cla()
    # for stopping simulation with the esc key.
    plt.gcf().canvas.mpl_connect('key_release_event',
            lambda event: [exit(0) if event.key == 'escape' else None])
    plt.plot(cx, cy, ".r", label="course")
    plt.plot(x, y, "-b", label="trajectory")
    plt.plot(cx[target_idx], cy[target_idx], "xg", label="target")
    plt.axis("equal")
    plt.grid(True)
    plt.title("Speed[km/h]:" + str(state.v * 3.6)[:4])
    plt.pause(0.001)

assert last_idx >= target_idx, "Cannot reach goal"

plt.plot(cx, cy, ".r", label="course")
plt.plot(x, y, "-b", label="trajectory")
plt.legend()
plt.xlabel("x[m]")
plt.ylabel("y[m]")
plt.axis("equal")
plt.grid(True)

plt.subplots(1)
plt.plot(t, [iv * 3.6 for iv in v], "-r")
plt.xlabel("Time[s]")
plt.ylabel("Speed[km/h]")
plt.grid(True)
plt.show()